package ru.ifmo.sound.neuro;

import ru.ifmo.sound.jmlbased.DatasetGenerator;
import ru.ifmo.sound.jmlbased.FeatureSelector;
import ru.ifmo.sound.jmlbased.SamplesLibrary;
import ru.ifmo.sound.jmlbased.WrapperMethods;

import javax.sound.sampled.UnsupportedAudioFileException;
import java.io.File;
import java.io.IOException;
import java.io.PrintStream;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * @author avhaliullin
 */
public class AllTheSystem {
    private Map<String, NeuroNet> nets = new HashMap<String, NeuroNet>();
    private Map<String, ExecutorService> executors = new HashMap<String, ExecutorService>();
    private PrintStream out = System.out;
    private FeatureSelector featureSelector;
    private SamplesLibrary samplesLibrary;
    private DatasetGenerator datasetGenerator;
    private Map<String, List<double[]>> currentDataset;
    private Float errorsOnLastRun = null;

    private void resetExecutors() {
        for (String className : samplesLibrary.listKnownClasses()) {
            executors.put(className, Executors.newFixedThreadPool(1));
        }
    }

    public float fullCycle() {
        resetExecutors();
        int errors = 0;
        int count = 0;
        if (errorsOnLastRun == null || errorsOnLastRun < 0.08) {
            System.out.println("Loading new dataset...");
            currentDataset = WrapperMethods.mapByDataset(featureSelector.createDataset(6));
        }
        List<Sample> samples = new ArrayList<Sample>();
        List<Executor> executorsList = new ArrayList<Executor>();
        for (String className : samplesLibrary.listKnownClasses()) {
            List<double[]> vectors = currentDataset.get(className);
            for (double[] vector : vectors) {
                samples.add(new Sample(className, vector));
            }
        }
        Collections.shuffle(samples);
        for (Sample sample : samples) {
            for (Map.Entry<String, NeuroNet> net : nets.entrySet()) {
                boolean correctResult = net.getKey().equals(sample.className);
                Executor executor = new Executor(net.getValue(), sample.vector, correctResult, net.getKey(), sample.className);
                executors.get(net.getKey()).execute(executor);
                executorsList.add(executor);

            }
        }
        for (ExecutorService executorService : executors.values()) {
            executorService.shutdown();
            try {
                executorService.awaitTermination(1, TimeUnit.MINUTES);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
        for (Executor executor : executorsList) {
            count++;
            errors += executor.wasError ? 1 : 0;
        }
        errorsOnLastRun = ((float) errors) / count;
        return errorsOnLastRun;
    }

    public Map<String, Boolean> analyze(double[] vector) {
        double[] filtered = featureSelector.filter(vector);
        Map<String, Boolean> result = new HashMap<String, Boolean>();
        for (Map.Entry<String, NeuroNet> net : nets.entrySet()) {
            result.put(net.getKey(), net.getValue().query(filtered));
        }
        System.out.println();

        return result;
    }

    public AllTheSystem(File workingDir) throws UnsupportedAudioFileException, IOException {
        samplesLibrary = new SamplesLibrary(workingDir);
        datasetGenerator = new DatasetGenerator(samplesLibrary);
        featureSelector = new FeatureSelector(datasetGenerator);
        for (String className : samplesLibrary.listKnownClasses()) {
            nets.put(className, NeuronetFactory.createSimple(FeatureSelector.inputs, 2, 0.000005));
        }
    }

    private class Sample {
        public final String className;
        public final double[] vector;

        Sample(String className, double[] vector) {
            this.className = className;
            this.vector = vector;
        }
    }

    class Executor implements Runnable {
        private final double[] vector;
        private final String name;
        private final String testName;
        private final NeuroNet net;
        private boolean wasError;
        private final boolean result;

        public Executor(NeuroNet net, double[] vector, boolean result, String name, String testName) {
            this.net = net;
            this.vector = vector;
            this.result = result;
            this.name = name;
            this.testName = testName;
        }

        public void run() {
            wasError = result ^ net.learn(vector, result);
//            synchronized (out) {
//                out.println(name + (wasError ? " failed " : " done ") + " with " + testName);
//            }
        }
    }
}
